"""
返回召回的结果
"""
from chatbot import config
from chatbot.dnn.recall.sentence_vectors import Sentence2Vector


class Recall:
    def __init__(self, by_word=False, method="bm25"):
        self.by_word = by_word
        self.sentence_vec = Sentence2Vector(by_word,method)
        self.vectorizer, self.features_vec, self.lines_cuted, self.search_index = self.sentence_vec.build_vectors()

    def predict(self, sentence):
        """
        :param sentence: {"cut_by_word":str,"cut":str,"entity":[str,str]}
        :return:
        """
        cur_sentence = [sentence["cut_by_word"]] if self.by_word else [sentence["cut"]]
        cur_sentence_vector = self.vectorizer.transform(cur_sentence)
        search_results = self.search_index.search(cur_sentence_vector, k=config.recall_topk,
                                                  k_clusters=config.recall_clusters, return_distance=True)
        # 过滤主体
        filted_results = []
        for result in search_results[0]:
            distance = result[0]
            key = result[1]
            entities = self.sentence_vec.qa_dict[key]["entity"]
            # 判断有交集
            if len(set(entities) & set(sentence["entity"])) > 0:
                filted_results.append(result)
        if len(filted_results) < 1:
            # 没交集
            return [i[1] for i in search_results[0]]
        else:
            return [i[1] for i in filted_results]
